package com.xzzz.common.lb;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.sql.Struct;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 自定义负载均衡策略: 根据CPU核数作为权重进行负载均衡
 * 在该体系中, CPU核数需要从服务实例中获取, 如果是 Nacos, 就需要在注册时携带本机的CPU核数来作为权重来计算.
 *
 * @author xzzz
 */
@Slf4j
public class CpuWeightLoadBalancer extends RoundRobinLoadBalancer {

    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public CpuWeightLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                 String serviceId) {
        super(serviceInstanceListSupplierProvider, serviceId);
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        log.info("负载均衡策略为: CPU权重");
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next()
                .map(serviceInstances -> processInstanceResponse(supplier, serviceInstances));
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            return new EmptyResponse();
        }

        List<Integer> weights = maintainWeights(instances);
        final int maxWeight = weights.size() == 0 ? 0 : weights.get(weights.size() - 1);
        // 随机权重
        double randomWeight = ThreadLocalRandom.current().nextInt(maxWeight);

        int n = 0;
        int index = 0;

        // 遍历权重, 如果当前机器的权重大于随机数, 则命中该机器
        // 本次的权重不需要像轮询一样依次递增, 在一定次数的请求后, 随机生成的权重与递增的权重几乎相同
        for (Integer weight : weights) {
            if (weight >= randomWeight) {
                index = n;
                break;
            } else {
                n++;
            }
        }

        ServiceInstance instance = instances.get(index);
        return new DefaultResponse(instance);
    }

    /**
     * 这是一种数值越大, 权重越大的算法
     *
     * @param instances 服务器列表
     * @return 权重集合
     */
    private List<Integer> maintainWeights(List<ServiceInstance> instances) {
        List<Integer> weights = new ArrayList<>();

        int weightSoFar = 0;

        for (ServiceInstance instance : instances) {
            // 核数越高, 权重越大, 集合的下标为 size()-1 即为目标服务集群的总核数
            final String core = instance.getMetadata().get("CPU-CORE");
            int weight = 1;
            if (StrUtil.isNotBlank(core)) {
                weight = Integer.parseInt(core);
            }
            weightSoFar += weight;
            weights.add(weightSoFar);
        }

        return weights;
    }

    /**
     * 这是一种数值越大, 权重越小的算法, 常用于通过响应时间, QPS来计算权重的场景, 如响应时间越大/QPS越高则该机器的权重越低
     *
     * @param instances 服务器列表
     * @return 权重集合
     */
    private List<Integer> maintainWeightsOther(List<ServiceInstance> instances) {
        List<Integer> weights = new ArrayList<>();

        // 总权重, 如果没有配置权重, 则总权重为0
        int totalWeight = instances.stream()
                .mapToInt(serviceInstance -> {
                    final String cpuCore = serviceInstance.getMetadata().get("CPU-CORE");
                    if (StrUtil.isBlank(cpuCore)) {
                        return 0;
                    }
                    return Integer.parseInt(cpuCore);
                })
                .sum();

        int weightSoFar = 0;

        for (ServiceInstance instance : instances) {
            // 核数越高, 权重越大, 集合的下标为 size()-1 即为目标服务集群的总核数
            final String core = instance.getMetadata().get("CPU-CORE");
            int weight = 1;
            if (StrUtil.isNotBlank(core)) {
                weight = totalWeight - Integer.parseInt(core);
            }
            weightSoFar += weight;
            weights.add(weightSoFar);
        }

        return weights;
    }
}
